from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv

class Hat:
    def __init__(self,systemRole=None, model="glm-4-plus"):
        self.model = model
        self.systemRole = systemRole
        _ = load_dotenv(find_dotenv())
        self.client = ZhipuAI()

    # 流式输出
    def stream_print(self, stream_response):
         # 逐步接收流式数据并显示
        resultContent = "" 
        for chunk in stream_response:
            hasChocies = hasattr(chunk, "choices") 
            if hasChocies and chunk.choices[0].finish_reason is None:
                content_chunk = chunk.choices[0].delta.content
                resultContent += content_chunk
                print(content_chunk, end="")
        print("\n")
        return resultContent

    # 与模型交互，生成一个想法
    def generate_idea(self, prompt, temperature=0.01):
        messages = []
        if self.systemRole is not None:
            messages.append({"role":"system", 
                             "content":self.systemRole})
        messages.append({"role":"user", "content":prompt}) 

        #print(f"\n\nprompt=\n{prompt}")
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature, 
            stream=True
        )
        resultContent = self.stream_print(response)
        return resultContent
    
    # 与模型交互，基于网络搜索生成一个想法
    def generate_info_from_Search(self, prompt, temperature=0.01):
        messages = []
        if self.systemRole is not None:
            messages.append({"role":"system", 
                             "content":self.systemRole})
        messages.append({"role":"user", "content":prompt}) 

        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,  
            stream=True,
            tool_choice="auto", 
            tools=[{
                "type": "web_search",
                "web_search": {
                    "enable": True,
                    "search_result": True,
                }
            }],
        )
        resultContent = self.stream_print(response)
        return resultContent

# 白帽
class WhiteHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "白帽"
        feature = "关注客观的事实和数据"
        function = "收集和分析与问题相关的客观信息，提供事实依据"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_info_from_Search(subject)

# 绿帽    
class GreenHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "绿帽"
        feature = "富有创造力和想象力"
        function = "提出新的想法、解决方案或可能性，鼓励创新和求异思维。"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_idea(subject)
    
# 黑帽
class BlackHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "黑帽"
        feature = "进行批判和质疑"
        function = "分析问题的缺点、风险和挑战，提出负面反馈和潜在问题。"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_idea(subject)
    
# 黄帽
class YellowHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "黄帽"
        feature = "从正面考虑问题，表达乐观的观点"
        function = "强调方案的优点和潜在好处，促进积极思考和建设性反馈。"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_idea(subject)
    
# 红帽
class RedHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "红帽"
        feature = "表达情感和直觉。"
        function = "分享个人感受和直觉，为决策提供情感层面的考虑。"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_idea(subject)
    
# 蓝帽
class BlueHat(Hat):
    def __init__(self,  *args, **kwargs):
        role = "蓝帽"
        feature = "负责控制和调节思维过程。"
        #function = "管理整个思考过程，确定讨论的目标、顺序和结论，确保所有观点都得到充分考虑。并负责做出结论"
        function = "确保所有观点都得到充分考虑，并负责根据其他角色的言论做出总结报告。"
        systemRole = f"""你的角色是六顶思考帽中的`{role}`，
            你的特点是`{feature}`。
            你的功能是`{function}`"""
        super().__init__(systemRole = systemRole, 
                         *args, **kwargs)

    def get_idea(self, subject):
        return super().generate_idea(subject)
    

def main():

    whitehat = WhiteHat()
    greenhat =GreenHat()
    yellowhat = YellowHat()
    blackhat = BlackHat()
    redhat = RedHat()
    bluehat = BlueHat()

    # 拟定一个讨论的话题
    topic = "氢能源汽车的前景"

    # 假设不同颜色的帽子发言顺序不定，以下执行顺序可随机调整；
    # 或者，可以线程或其它形式并发执行
    print("【白帽】")
    idea_white = whitehat.get_idea(topic)
    print("【绿帽】")
    idea_green = greenhat.get_idea(topic)
    print("【黑帽】")
    idea_black = blackhat.get_idea(topic)
    print("【黄帽】")
    idea_yellow = yellowhat.get_idea(topic)
    print("【红帽】")
    idea_red = redhat.get_idea(topic)

    # 蓝帽总结各方发言
    idea_all = f"{idea_white}{idea_green}{idea_black}{idea_yellow}{idea_red}"
    print("【蓝帽】")
    idea_blue = bluehat.get_idea(idea_all)

    # 输出蓝帽对各方发言的总结
    # print(f"【蓝帽】\n{idea_blue}")

if __name__ == "__main__":
    main()